Skip to content

Fix alias output naming and improve Spike I/O diagnostics#51

Merged
hgt312 merged 1 commit into
aws-neuron:mainfrom
hgt312:alias-output-naming-fix
Jul 15, 2026
Merged

Fix alias output naming and improve Spike I/O diagnostics#51
hgt312 merged 1 commit into
aws-neuron:mainfrom
hgt312:alias-output-naming-fix

Conversation

@hgt312

@hgt312 hgt312 commented Mar 31, 2026

Copy link
Copy Markdown
Contributor

Fix alias output naming and improve Spike I/O diagnostics

Problem

The tracer's output naming step (Step 3 in NKIPyKernel._build_code) used a truthiness check (if not r.backend_tensor.name) to decide whether to assign the canonical
output{idx} name. This failed when a non-aliased output tensor already carried a name from tracing (e.g. "intermediate0" from an np.add result). The tensor kept
its intermediate name, causing a mismatch between the NEFF I/O table and what callers pass to kernel(inputs={...}, outputs={...}).

This was discovered during sglang-nkipy integration with the stable neuronx-cc 2.23.6484.0 compiler, where three kernel call patterns broke:

  • Standalone aliased kernels (e.g. update_kv_cache): NEFF input renamed to kv_cache.must_alias_input but callers still passed kv_cache
  • Fused graphs with aliased params (e.g. prefill_pre_moe): aliased output named kv_cache shifted other outputs to output1, output2, ...
  • Broken-identity aliases (e.g. prefill_post_moe): mutated param passed through an NKI wrapper that breaks tensor identity, causing the tracer to auto-append a 3rd
    output

Debugging these required extracting NEFF I/O names from HLO protobuf binaries because Spike's _validate_io only produced a bare KeyError.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

@hgt312
hgt312 requested a review from a team March 31, 2026 17:13
Comment thread spike/src/spike/spike_model.py Outdated
model_core_id = self.model_ref.core_id

unknown_inputs = set(inputs) - set(self.input_tensors_info)
if unknown_inputs:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like the checks!

Always lower non-aliased results as output{idx}, even when tracing assigned an intermediate name. Keep alias naming in NKIPy lowering, add direct NEFF naming coverage, and improve Spike I/O mismatch diagnostics without teaching Spike alias policy.
@hgt312
hgt312 force-pushed the alias-output-naming-fix branch from aa46211 to 41bd1a8 Compare July 10, 2026 06:44
@hgt312
hgt312 requested a review from vgene July 10, 2026 23:58
@hgt312

hgt312 commented Jul 13, 2026

Copy link
Copy Markdown
Contributor Author

@vgene can you help review again?

@vgene vgene left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@ymwangg
ymwangg requested a review from a team July 14, 2026 16:37

@ymwangg ymwangg left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@hgt312
hgt312 merged commit 17365ae into aws-neuron:main Jul 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants